from copy import deepcopy
from random import randint

from utils.util import Utils
from utils.generator import Generator
from utils.profile import Profile


class NetherlandFlag(Utils):
    # ref: https://www.bilibili.com/video/BV1cc411F7Y6/?spm_id_from=autoNext&vd_source=ad559e41b561d48861f70d0187e51e02
    # 27:00
    def __init__(self, arr):
        super(NetherlandFlag, self).__init__(arr)
        # self.arr_backup = deepcopy(arr)

    @Profile()
    def partition(self, left, right, x):
        first = left
        last = right
        i = left
        while i <= last:
            if self.arr[i] == x:
                i += 1
            elif self.arr[i] < x:
                self.swap(first, i)
                first += 1
                i += 1
            else:
                self.swap(i, last)
                last -= 1
        return first, last

    @Profile()
    def partition_optimize(self, left, right, x):
        first = left
        last = right
        i = left
        while self.arr[last] > x:
            last -= 1
        while i <= last:
            if self.arr[i] == x:
                i += 1
            elif self.arr[i] < x:
                self.swap(first, i)
                first += 1
                i += 1
            else:
                self.swap(i, last)
                last -= 1
                while self.arr[last] > x:
                    last -= 1
        return first, last


class FastSort(NetherlandFlag):
    def __init__(self, arr):
        super(FastSort, self).__init__(arr)

    @Profile()
    def sort(self, left=None, right=None):
        left = 0 if left is None else left
        right = self.len - 1 if right is None else right
        if left >= right:
            return
        x = self.arr[randint(left, right)]
        first, last = self.partition_optimize(left, right, x)
        self.sort(left, first - 1)
        self.sort(last + 1, right)

    @Profile()
    def classic_sort(self, left=None, right=None):
        left = 0 if left is None else left
        right = self.len - 1 if right is None else right
        if left >= right:
            return
        x = self.arr[randint(left, right)]
        middle = self.classic_partition(left, right, x)
        self.classic_sort(left, middle - 1)
        self.classic_sort(middle + 1, right)

    @Profile()
    def classic_partition(self, left, right, x):
        gt_index = left
        x_index = left
        i = left
        while i <= right:
            if self.arr[i] <= x:
                self.swap(gt_index, i)
                if self.arr[gt_index] == x:
                    x_index = gt_index
                gt_index += 1
            i += 1
        self.swap(x_index, gt_index - 1)
        return gt_index - 1


if __name__ == '__main__':
    arr = Generator(min_val=10, max_val=30).do()
    arr2 = deepcopy(arr)
    print(arr)
    FastSort(arr).classic_sort()
    Profile().print_and_reset()

    FastSort(arr2).sort()
    print(arr)
    print(arr2)
    print(arr == arr2)
